import os
import csv
from ..build_embeddings import EmbeddingBuilder
from ..io.omics_dataset_writer import OmicsDatasetWriter
from zkl_pyutils_fsspec import resolve_fs
from zkl_pyutils_promptui import confirm_overwrite_path

class EncodingEmbeddingWriter:
    def __init__(self, path, model_name="Qwen/Qwen3-Embedding-0.6B", tensor_parallel_size=2):
        # 确认目标目录是否清空
        confirm_overwrite_path(path)

        self.builder = EmbeddingBuilder(model_name, tensor_parallel_size=tensor_parallel_size)
        self.path = path
        self.dataset_path = os.path.join(self.path, "dataset")

        if not os.path.exists(self.path):
            os.makedirs(self.path)

        self.omics_dataset_init()
        self.csv_init()

    def csv_init(self):
        """Initialize CSV file for writing encoding keys."""
        self.csv_file_path = os.path.join(self.path, "encoding_keys.csv")
        self.csv_file = open(self.csv_file_path, 'w', newline='', encoding='utf-8')
        self.csv_writer = csv.writer(self.csv_file)
        self.csv_writer.writerow(['key'])

    def omics_dataset_init(self):
        """Initialize OmicsDataset for writing embeddings."""
        # confirm_overwrite_path(self.path)
        dst = resolve_fs(self.dataset_path)
        self.dataset_writer = OmicsDatasetWriter(dst)
        self.str_embeddings_writer = self.dataset_writer.str_embeddings_writer()
        self.str_embeddings_writer.write_columns_name(['embeddings'])

    def embedding_and_write_batch(self, encoding_items):
        """Embeds a batch of texts and writes embeddings and keys."""
        texts = [item[1] for item in encoding_items]
        keys = [item[0] for item in encoding_items]

        embeddings = self.builder.embed(texts)

        for key, embedding in zip(keys, embeddings):
            # print(key , embedding)
            self.csv_writer.writerow([key])
            self.str_embeddings_writer.write([embedding])

    def close(self):
        self.dataset_writer.close()
        if hasattr(self, 'csv_file') and self.csv_file and not self.csv_file.closed:
            self.csv_file.close()

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()